{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "597c13c0",
   "metadata": {},
   "source": [
    "# Serving a Stable Diffusion Model with Ray Serve\n",
    "\n",
    "| Template Specification | Description |\n",
    "| ---------------------- | ----------- |\n",
    "| Summary | This template loads a pretrained stable diffusion model from HuggingFace and serves it to a local endpoint as a [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) deployment. |\n",
    "| Time to Run | Around 2 minutes to setup the models and generate your first image(s). Less than 10 seconds for every subsequent round of image generation (depending on the image size). |\n",
    "| Minimum Compute Requirements | At least 1 GPU node. The default is 4 nodes, each with 1 NVIDIA T4 GPU. |\n",
    "| Cluster Environment | This template uses a custom docker image built on top of the Anyscale-provided Ray image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview). See the appendix in the `README` for more details. |\n",
    "\n",
    "By the end, we'll have an application that generates images using stable diffusion for a given prompt!\n",
    "\n",
    "The application will look something like this:\n",
    "\n",
    "```text\n",
    "Enter a prompt (or 'q' to quit):   twin peaks sf in basquiat painting style\n",
    "\n",
    "Generating image(s)...\n",
    "\n",
    "Generated 4 image(s) in 8.75 seconds to the directory: 58b298d9\n",
    "```\n",
    "\n",
    "![Example output](https://github-production-user-asset-6210df.s3.amazonaws.com/3887863/239090189-dc1f1b7b-2fa0-4886-ae12-ca5d35b8ebc9.png)\n",
    "> Slot in your code below wherever you see the ✂️ icon to build off of this template!\n",
    ">\n",
    "> The framework and data format used in this template can be easily replaced to suit your own application!\n",
    "\n",
    "We'll start with some imports and initialize Ray:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59842da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastapi import FastAPI\n",
    "from fastapi.responses import Response\n",
    "from io import BytesIO\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import requests\n",
    "import time\n",
    "import uuid\n",
    "\n",
    "import ray\n",
    "from ray import serve\n",
    "\n",
    "ray.init()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "520ef4d7",
   "metadata": {},
   "source": [
    "## Deploy the Ray Serve application locally\n",
    "\n",
    "First, we define the Ray Serve application with the model loading and inference logic. This includes setting up:\n",
    "- The `/imagine` API endpoint that we query to generate the image.\n",
    "- The stable diffusion model loaded inside a Ray Serve Deployment.\n",
    "  We'll specify the *number of model replicas* to keep active in our Ray cluster. These model replicas can process incoming requests concurrently.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "de6318ac",
   "metadata": {},
   "source": [
    "> ✂️ Replace these values to change the number of model replicas to serve, as well as the GPU resources required by each replica.\n",
    ">\n",
    "> With more model replicas, more images can be generated in parallel!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90eca147",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "NUM_REPLICAS: int = 4\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be41ca9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "if NUM_REPLICAS > ray.available_resources()[\"GPU\"]:\n",
    "    print(\n",
    "        \"Your cluster does not currently have enough resources to run with these settings. \"\n",
    "        \"Consider decreasing the number of workers, or decreasing the resources needed \"\n",
    "        \"per worker. Ignore this if your cluster auto-scales.\"\n",
    "    )\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "89eb3e2c",
   "metadata": {},
   "source": [
    "First, we define the Ray Serve Deployment, which will load a stable diffusion model and perform inference with it.\n",
    "\n",
    "> ✂️ Modify this block to load your own model, and change the `generate` method to perform your own online inference logic!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f203efd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@serve.deployment(\n",
    "    ray_actor_options={\"num_gpus\": 1},\n",
    "    num_replicas=NUM_REPLICAS,\n",
    ")\n",
    "class StableDiffusionV2:\n",
    "    def __init__(self):\n",
    "        # <Replace with your own model loading logic>\n",
    "        import torch\n",
    "        from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline\n",
    "\n",
    "        model_id = \"stabilityai/stable-diffusion-2\"\n",
    "        scheduler = EulerDiscreteScheduler.from_pretrained(\n",
    "            model_id, subfolder=\"scheduler\"\n",
    "        )\n",
    "        self.pipe = StableDiffusionPipeline.from_pretrained(\n",
    "            model_id, scheduler=scheduler, revision=\"fp16\", torch_dtype=torch.float16\n",
    "        )\n",
    "        self.pipe = self.pipe.to(\"cuda\")\n",
    "\n",
    "    def generate(self, prompt: str, img_size: int = 776):\n",
    "        # <Replace with your own model inference logic>\n",
    "        assert len(prompt), \"prompt parameter cannot be empty\"\n",
    "        image = self.pipe(prompt, height=img_size, width=img_size).images[0]\n",
    "        return image\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0134aa54",
   "metadata": {},
   "source": [
    "Next, we'll define the actual API endpoint to live at `/imagine`.\n",
    "\n",
    "> ✂️ Modify this block to change the endpoint URL, response schema, and add any post-processing logic needed from your model output!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f80fee2",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = FastAPI()\n",
    "\n",
    "\n",
    "@serve.deployment(num_replicas=1)\n",
    "@serve.ingress(app)\n",
    "class APIIngress:\n",
    "    def __init__(self, diffusion_model_handle) -> None:\n",
    "        self.handle = diffusion_model_handle\n",
    "\n",
    "    @app.get(\n",
    "        \"/imagine\",\n",
    "        responses={200: {\"content\": {\"image/png\": {}}}},\n",
    "        response_class=Response,\n",
    "    )\n",
    "    async def generate(self, prompt: str, img_size: int = 776):\n",
    "        assert len(prompt), \"prompt parameter cannot be empty\"\n",
    "\n",
    "        image = await self.handle.generate.remote(prompt, img_size=img_size)\n",
    "\n",
    "        file_stream = BytesIO()\n",
    "        image.save(file_stream, \"PNG\")\n",
    "        return Response(content=file_stream.getvalue(), media_type=\"image/png\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61b8916d",
   "metadata": {},
   "source": [
    "Now, we deploy the Ray Serve application locally at `http://localhost:8000`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfc2e244",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "entrypoint = APIIngress.bind(StableDiffusionV2.bind())\n",
    "\n",
    "# Shutdown any existing Serve replicas, if they're still around.\n",
    "serve.shutdown()\n",
    "serve.run(entrypoint, name=\"serving_stable_diffusion_template\")\n",
    "print(\"Done setting up replicas! Now accepting requests...\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "757678cc",
   "metadata": {},
   "source": [
    "## Make requests to the endpoint\n",
    "\n",
    "Next, we'll build a simple client to submit prompts as HTTP requests to the local endpoint at `http://localhost:8000/imagine`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "008976b5",
   "metadata": {},
   "source": [
    "Start the client script in the next few cells, and generate your first image!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67ad095b",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = \"http://localhost:8000/imagine\"\n",
    "\n",
    "\n",
    "@ray.remote(num_cpus=0)\n",
    "def generate_image(prompt, image_size):\n",
    "    req = {\"prompt\": prompt, \"img_size\": image_size}\n",
    "    resp = requests.get(endpoint, params=req)\n",
    "    return resp.content\n",
    "\n",
    "\n",
    "def show_images(filenames):\n",
    "    fig, axs = plt.subplots(1, len(filenames), figsize=(4 * len(filenames), 4))\n",
    "    for i, filename in enumerate(filenames):\n",
    "        ax = axs if len(filenames) == 1 else axs[i]\n",
    "        ax.imshow(plt.imread(filename))\n",
    "        ax.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def main(\n",
    "    interactive: bool = False,\n",
    "    prompt: str = \"twin peaks sf in basquiat painting style\",\n",
    "    num_images: int = 4,\n",
    "    image_size: int = 640,\n",
    "):\n",
    "    try:\n",
    "        requests.get(endpoint, timeout=0.1)\n",
    "    except Exception as e:\n",
    "        raise RuntimeWarning(\n",
    "            \"Did you setup the Ray Serve model replicas with `serve.run` \"\n",
    "            \"in a previous cell?\"\n",
    "        ) from e\n",
    "\n",
    "    generation_times = []\n",
    "    while True:\n",
    "        prompt = (\n",
    "            prompt\n",
    "            if not interactive\n",
    "            else input(f\"\\nEnter a prompt (or 'q' to quit):  \")\n",
    "        )\n",
    "        if prompt.lower() == \"q\":\n",
    "            break\n",
    "\n",
    "        print(\"\\nGenerating image(s)...\\n\")\n",
    "        start = time.time()\n",
    "\n",
    "        # Make `num_images` requests to the endpoint at once!\n",
    "        images = ray.get(\n",
    "            [generate_image.remote(prompt, image_size) for _ in range(num_images)]\n",
    "        )\n",
    "\n",
    "        dirname = f\"{uuid.uuid4().hex[:8]}\"\n",
    "        os.makedirs(dirname)\n",
    "        filenames = []\n",
    "        for i, image in enumerate(images):\n",
    "            filename = os.path.join(dirname, f\"{i}.png\")\n",
    "            with open(filename, \"wb\") as f:\n",
    "                f.write(image)\n",
    "            filenames.append(filename)\n",
    "\n",
    "        elapsed = time.time() - start\n",
    "        generation_times.append(elapsed)\n",
    "        print(\n",
    "            f\"\\nGenerated {len(images)} image(s) in {elapsed:.2f} seconds to \"\n",
    "            f\"the directory: {dirname}\\n\"\n",
    "        )\n",
    "        show_images(filenames)\n",
    "        if not interactive:\n",
    "            break\n",
    "    return np.mean(generation_times) if generation_times else -1\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c8949cc7",
   "metadata": {},
   "source": [
    "Once the stable diffusion model finishes generating your image(s), it will be included in the HTTP response body.\n",
    "The client saves all the images in a local directory for you to view, and they'll also show up in the notebook cell!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3e29193b",
   "metadata": {},
   "source": [
    "> ✂️ Replace this value to change the number of images to generate per prompt.\n",
    ">\n",
    "> Each image will be generated starting from a different set of random noise,\n",
    "> so you'll be able to see multiple options per prompt!\n",
    ">\n",
    "> Try starting with `NUM_IMAGES_PER_PROMPT` equal to `NUM_REPLICAS` from earlier.\n",
    ">\n",
    "> You can choose to run this interactively, or submit a single `PROMPT`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd20a52d",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_IMAGES_PER_PROMPT: int = NUM_REPLICAS\n",
    "\n",
    "# Control the output size: (IMAGE_SIZE, IMAGE_SIZE)\n",
    "# The stable diffusion model requires `IMAGE_SIZE` to be a multiple of 8.\n",
    "# NOTE: Generated image quality degrades rapidly if you reduce the size too much.\n",
    "IMAGE_SIZE: int = 640\n",
    "\n",
    "INTERACTIVE: bool = False\n",
    "PROMPT = \"twin peaks sf in basquiat painting style\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_generation_time = main(\n",
    "    interactive=INTERACTIVE,\n",
    "    prompt=PROMPT,\n",
    "    num_images=NUM_IMAGES_PER_PROMPT,\n",
    "    image_size=IMAGE_SIZE,\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fb124968",
   "metadata": {},
   "source": [
    "You've successfully served a stable diffusion model!\n",
    "You can modify this template and iterate your model deployment directly on your cluster within your Anyscale Workspace,\n",
    "testing with the local endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e360cf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Shut down the model replicas once you're done!\n",
    "serve.shutdown()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "880c2d6f",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This template used [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) to serve many replicas of a stable diffusion model. \n",
    "\n",
    "At a high level, this template showed how to:\n",
    "1. Define a Ray Serve deployment to load a HuggingFace model and perform inference.\n",
    "2. Set up a local endpoint to accept and route requests to the different model replicas.\n",
    "3. Make multiple requests in parallel to generate many images at a time.\n",
    "\n",
    "See this [getting started guide](https://docs.ray.io/en/latest/serve/getting_started.html) for a more detailed walkthrough of Ray Serve."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bcc69b2d",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ray_dev_py38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "265d195fda5292fe8f69c6e37c435a5634a1ed3b6799724e66a975f68fa21517"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
